//! A GPU accelerated JPEG codec library.
//!
//! # Examples
//!
//! ```
//! use nvjpeg::Decoder;
/// ```
use std::convert::TryInto;
use std::ptr::NonNull;

pub use nvjpeg_sys as ffi;

pub mod decoder;
pub use decoder::*;

pub const MAX_COMPONENT: usize = ffi::NVJPEG_MAX_COMPONENT as usize;

/// The Result for this crate.
pub type Result<T> = std::result::Result<T, Error>;

/// The nvJPEG chroma subsampling.
#[repr(u32)]
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum ChromaSubsampling {
    Css444,
    Css422,
    Css420,
    Css440,
    Css411,
    Css410,
    CssGray,
    #[cfg(feature = "cuda-11-4")]
    Css410V,
    CssUnknown,
}

impl Default for ChromaSubsampling {
    fn default() -> Self {
        Self::CssUnknown
    }
}

impl From<ffi::nvjpegChromaSubsampling_t> for ChromaSubsampling {
    fn from(val: ffi::nvjpegChromaSubsampling_t) -> Self {
        use ffi::nvjpegChromaSubsampling_t::*;
        use ChromaSubsampling::*;

        match val {
            NVJPEG_CSS_444 => Css444,
            NVJPEG_CSS_422 => Css422,
            NVJPEG_CSS_420 => Css420,
            NVJPEG_CSS_440 => Css440,
            NVJPEG_CSS_411 => Css411,
            NVJPEG_CSS_410 => Css410,
            NVJPEG_CSS_GRAY => CssGray,
            #[cfg(feature = "cuda-11-4")]
            NVJPEG_CSS_410V => Css410V,
            NVJPEG_CSS_UNKNOWN => CssUnknown,
        }
    }
}

impl From<ChromaSubsampling> for ffi::nvjpegChromaSubsampling_t {
    fn from(val: ChromaSubsampling) -> Self {
        use ffi::nvjpegChromaSubsampling_t::*;
        use ChromaSubsampling::*;

        match val {
            Css444 => NVJPEG_CSS_444,
            Css422 => NVJPEG_CSS_422,
            Css420 => NVJPEG_CSS_420,
            Css440 => NVJPEG_CSS_440,
            Css411 => NVJPEG_CSS_411,
            Css410 => NVJPEG_CSS_410,
            CssGray => NVJPEG_CSS_GRAY,
            #[cfg(feature = "cuda-11-4")]
            Css410V => NVJPEG_CSS_410V,
            CssUnknown => NVJPEG_CSS_UNKNOWN,
        }
    }
}

/// A memory block allocated for cuda host.
#[derive(Debug, PartialEq)]
pub struct CudaHostMemory {
    ptr: NonNull<u8>,
    size: usize,
}

impl CudaHostMemory {
    pub fn new(size: usize) -> Result<Self> {
        use ffi::cudaError::cudaSuccess;

        unsafe {
            let mut ptr: *mut std::os::raw::c_void = std::ptr::null_mut();
            match ffi::cudaMallocHost(&mut ptr, size.try_into().unwrap()) {
                cudaSuccess => Ok(Self {
                    ptr: NonNull::new_unchecked(ptr as *mut u8),
                    size,
                }),
                err => Err(Error::Cuda(err)),
            }
        }
    }

    pub fn as_bytes(&self) -> &[u8] {
        unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.size) }
    }

    pub fn as_bytes_mut(&mut self) -> &mut [u8] {
        unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.size) }
    }

    pub fn as_ptr(&self) -> *const u8 {
        self.ptr.as_ptr()
    }

    pub fn as_mut_ptr(&self) -> *mut u8 {
        self.ptr.as_ptr()
    }

    pub fn is_empty(&self) -> bool {
        self.size == 0
    }

    pub fn len(&self) -> usize {
        self.size
    }
}

impl Drop for CudaHostMemory {
    fn drop(&mut self) {
        unsafe {
            ffi::cudaFreeHost(self.ptr.as_ptr() as *mut std::os::raw::c_void);
        }
    }
}

/// A cuda stream handle.
pub struct CudaStream {
    ptr: NonNull<ffi::CUstream_st>,
}

impl CudaStream {
    pub fn new() -> Result<Self> {
        use ffi::cudaError::cudaSuccess;

        unsafe {
            let mut ptr: ffi::cudaStream_t = std::ptr::null_mut();
            match ffi::cudaStreamCreate(&mut ptr) {
                cudaSuccess => Ok(Self {
                    ptr: NonNull::new_unchecked(ptr),
                }),
                err => Err(Error::Cuda(err)),
            }
        }
    }

    pub fn as_ptr(&self) -> *mut ffi::CUstream_st {
        self.ptr.as_ptr()
    }
}

impl Drop for CudaStream {
    fn drop(&mut self) {
        unsafe {
            ffi::cudaStreamDestroy(self.ptr.as_ptr());
        }
    }
}

/// A type represents the library errors.
#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Error {
    Cuda(ffi::cudaError),
    Jpeg(ffi::nvjpegStatus_t),
}

/// A nvJPEG library handle.
#[derive(Debug)]
pub struct Handle {
    ptr: NonNull<ffi::nvjpegHandle>,
}

impl Handle {
    pub fn new() -> Result<Self> {
        use ffi::nvjpegStatus_t::NVJPEG_STATUS_SUCCESS;

        unsafe {
            let mut handle: ffi::nvjpegHandle_t = std::ptr::null_mut();
            match ffi::nvjpegCreateSimple(&mut handle) {
                NVJPEG_STATUS_SUCCESS => Ok(Self {
                    ptr: NonNull::new_unchecked(handle),
                }),
                status => Err(Error::Jpeg(status)),
            }
        }
    }

    pub fn as_ptr(&self) -> *mut ffi::nvjpegHandle {
        self.ptr.as_ptr()
    }
}

impl Drop for Handle {
    fn drop(&mut self) {
        unsafe {
            ffi::nvjpegDestroy(self.ptr.as_ptr());
        }
    }
}

/// A decoded image.
#[derive(Debug, PartialEq)]
pub struct Image {
    plane: [*mut u8; MAX_COMPONENT],
    pitch: [usize; MAX_COMPONENT],
    info: ImageInfo,
    pix_fmt: PixelFormat,
    mem: CudaHostMemory,
}

impl Image {
    pub fn new(info: ImageInfo, pix_fmt: PixelFormat) -> Result<Self> {
        use PixelFormat::*;

        let width = info.widths[0];
        let height = info.heights[0];

        match (info.components, pix_fmt) {
            (3, RgbPlanar | BgrPlanar) => {
                let picth = width;
                let offset = (picth * height) as isize;
                let size = picth * height * 3;
                let mem = CudaHostMemory::new(size)?;
                let plane0 = mem.as_mut_ptr();
                let plane1 = unsafe { plane0.offset(offset) };
                let plane2 = unsafe { plane1.offset(offset) };
                Ok(Self {
                    plane: [plane0, plane1, plane2, std::ptr::null_mut()],
                    pitch: [picth, picth, picth, 0],
                    info,
                    pix_fmt,
                    mem,
                })
            }
            (3, RgbInterleaved | BgrInterleaved) => {
                let picth = width * 3;
                let size = picth * height;
                let mem = CudaHostMemory::new(size)?;
                Ok(Self {
                    plane: [
                        mem.as_mut_ptr(),
                        std::ptr::null_mut(),
                        std::ptr::null_mut(),
                        std::ptr::null_mut(),
                    ],
                    pitch: [picth, 0, 0, 0],
                    info,
                    pix_fmt,
                    mem,
                })
            }
            _ => todo!(),
        }
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.mem.as_bytes()
    }

    pub fn as_bytes_mut(&mut self) -> &mut [u8] {
        self.mem.as_bytes_mut()
    }

    pub fn as_ptr(&self) -> *const u8 {
        self.mem.as_ptr()
    }

    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        self.mem.as_mut_ptr()
    }

    pub fn components(&self) -> usize {
        self.info.components
    }

    pub fn dimensions(&self) -> (usize, usize) {
        (self.width(), self.height())
    }

    pub fn width(&self) -> usize {
        self.info.widths[0]
    }

    pub fn width_of(&self, index: usize) -> usize {
        self.info.widths[index]
    }

    pub fn height(&self) -> usize {
        self.info.heights[0]
    }

    pub fn height_of(&self, index: usize) -> usize {
        self.info.heights[index]
    }

    pub fn pixel_format(&self) -> PixelFormat {
        self.pix_fmt
    }

    pub fn memory(&self) -> &CudaHostMemory {
        &self.mem
    }

    pub fn pitch_of(&self, index: usize) -> usize {
        self.pitch[index]
    }

    pub fn pitchs(&self) -> &[usize] {
        &self.pitch
    }

    pub fn plane_bytes_of(&self, index: usize) -> &[u8] {
        unsafe { std::slice::from_raw_parts(self.plane[index], self.info.heights[index]) }
    }

    pub fn plane_bytes_mut_of(&self, index: usize) -> &[u8] {
        unsafe { std::slice::from_raw_parts_mut(self.plane[index], self.info.heights[index]) }
    }

    pub fn planes(&self) -> &[*mut u8] {
        &self.plane
    }

    pub fn to_nvjpeg_image(&self) -> ffi::nvjpegImage_t {
        ffi::nvjpegImage_t {
            channel: [self.plane[0], self.plane[1], self.plane[2], self.plane[3]],
            pitch: [
                self.pitch[0].try_into().unwrap(),
                self.pitch[1].try_into().unwrap(),
                self.pitch[2].try_into().unwrap(),
                self.pitch[3].try_into().unwrap(),
            ],
        }
    }
}

/// The basic information about the image.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct ImageInfo {
    /// The number of components.
    pub components: usize,
    /// The chroma subsampling for the 1- or 3- channel encoding.
    pub subsampling: ChromaSubsampling,
    /// the width of each channel, in pixels.
    pub widths: [usize; MAX_COMPONENT],
    /// The height of each channel, in pixels.
    pub heights: [usize; MAX_COMPONENT],
}

/// A decoded image (readonly).
#[derive(Debug)]
pub struct ImageView(Image);

unsafe impl Send for ImageView {}
unsafe impl Sync for ImageView {}

impl ImageView {
    pub fn new(image: Image) -> Self {
        Self(image)
    }

    pub fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }

    pub fn as_ptr(&self) -> *const u8 {
        self.0.as_ptr()
    }

    pub fn components(&self) -> usize {
        self.0.components()
    }

    pub fn dimensions(&self) -> (usize, usize) {
        (self.0.width(), self.0.height())
    }

    pub fn width(&self) -> usize {
        self.0.width()
    }

    pub fn height(&self) -> usize {
        self.0.width()
    }

    pub fn height_of(&self, index: usize) -> usize {
        self.0.height_of(index)
    }

    pub fn pixel_format(&self) -> PixelFormat {
        self.0.pixel_format()
    }

    pub fn pitch_of(&self, index: usize) -> usize {
        self.0.pitch_of(index)
    }

    pub fn pitchs(&self) -> &[usize] {
        self.0.pitchs()
    }

    pub fn plane_bytes_of(&self, index: usize) -> &[u8] {
        self.0.plane_bytes_of(index)
    }

    pub fn planes(&self) -> &[*mut u8] {
        self.0.planes()
    }

    pub fn into_inner(self) -> Image {
        self.0
    }
}

/// A jpeg decoding temporary state handle.
pub struct JpegState {
    ptr: NonNull<ffi::nvjpegJpegState>,
}

impl JpegState {
    pub fn new(handle: &Handle) -> Result<Self> {
        use ffi::nvjpegStatus_t::NVJPEG_STATUS_SUCCESS;

        unsafe {
            let mut ptr: ffi::nvjpegJpegState_t = std::ptr::null_mut();
            match ffi::nvjpegJpegStateCreate(handle.as_ptr(), &mut ptr) {
                NVJPEG_STATUS_SUCCESS => Ok(Self {
                    ptr: NonNull::new_unchecked(ptr),
                }),
                status => Err(Error::Jpeg(status)),
            }
        }
    }

    pub fn as_ptr(&self) -> *mut ffi::nvjpegJpegState {
        self.ptr.as_ptr()
    }
}

impl Drop for JpegState {
    fn drop(&mut self) {
        unsafe {
            ffi::nvjpegJpegStateDestroy(self.ptr.as_ptr());
        }
    }
}

/// An enumeration in which the decoded output will be saved.
#[repr(u32)]
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum PixelFormat {
    /// Return the decoded image planar format.
    Unchanged,
    /// Convert to planar RGB.
    RgbPlanar,
    /// Convert to planar BGR.
    BgrPlanar,
    /// Convert to interleaved RGB.
    RgbInterleaved,
    /// Convert to interleaved BGR.
    BgrInterleaved,
    /// Return the Y component only.
    Y,
    /// Return in the YUV planar format.
    Yuv,
}

impl From<ffi::nvjpegOutputFormat_t> for PixelFormat {
    fn from(val: ffi::nvjpegOutputFormat_t) -> Self {
        use ffi::nvjpegOutputFormat_t::*;
        use PixelFormat::*;

        match val {
            NVJPEG_OUTPUT_UNCHANGED => Unchanged,
            NVJPEG_OUTPUT_RGB => RgbPlanar,
            NVJPEG_OUTPUT_BGR => BgrPlanar,
            NVJPEG_OUTPUT_RGBI => RgbInterleaved,
            NVJPEG_OUTPUT_BGRI => BgrInterleaved,
            NVJPEG_OUTPUT_Y => Y,
            NVJPEG_OUTPUT_YUV => Yuv,
        }
    }
}

impl From<PixelFormat> for ffi::nvjpegOutputFormat_t {
    fn from(val: PixelFormat) -> Self {
        use ffi::nvjpegOutputFormat_t::*;
        use PixelFormat::*;

        match val {
            Unchanged => NVJPEG_OUTPUT_UNCHANGED,
            RgbPlanar => NVJPEG_OUTPUT_RGB,
            BgrPlanar => NVJPEG_OUTPUT_BGR,
            RgbInterleaved => NVJPEG_OUTPUT_RGBI,
            BgrInterleaved => NVJPEG_OUTPUT_BGRI,
            Y => NVJPEG_OUTPUT_Y,
            Yuv => NVJPEG_OUTPUT_YUV,
        }
    }
}

/// A memory block allocated on system heap.
#[derive(Debug, PartialEq)]
pub struct SysMemory {
    ptr: NonNull<u8>,
    layout: std::alloc::Layout,
}

impl SysMemory {
    pub fn new(size: usize) -> Result<Self> {
        unsafe {
            let layout = std::alloc::Layout::from_size_align_unchecked(size, 4096);
            let ptr = std::alloc::alloc(layout);
            // assert_eq!(ptr, std::ptr::null_mut());
            Ok(Self {
                ptr: NonNull::new_unchecked(ptr),
                layout,
            })
        }
    }

    pub fn as_bytes(&self) -> &[u8] {
        unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.layout.size()) }
    }

    pub fn as_bytes_mut(&mut self) -> &mut [u8] {
        unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.layout.size()) }
    }

    pub fn as_ptr(&self) -> *const u8 {
        self.ptr.as_ptr()
    }

    pub fn as_mut_ptr(&self) -> *mut u8 {
        self.ptr.as_ptr()
    }

    pub fn is_empty(&self) -> bool {
        self.layout.size() == 0
    }

    pub fn len(&self) -> usize {
        self.layout.size()
    }
}

impl Drop for SysMemory {
    fn drop(&mut self) {
        unsafe {
            std::alloc::dealloc(self.ptr.as_ptr(), self.layout);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_decoder() {
        const LENA_JPG: &[u8] = &[
            0xff, 0xd8, 0xff, 0xe0, 0x00, 0x10, 0x4a, 0x46, 0x49, 0x46, 0x00, 0x01, 0x01, 0x00,
            0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0xff, 0xdb, 0x00, 0x43, 0x00, 0x10, 0x0b, 0x0c,
            0x0e, 0x0c, 0x0a, 0x10, 0x0e, 0x0d, 0x0e, 0x12, 0x11, 0x10, 0x13, 0x18, 0x28, 0x1a,
            0x18, 0x16, 0x16, 0x18, 0x31, 0x23, 0x25, 0x1d, 0x28, 0x3a, 0x33, 0x3d, 0x3c, 0x39,
            0x33, 0x38, 0x37, 0x40, 0x48, 0x5c, 0x4e, 0x40, 0x44, 0x57, 0x45, 0x37, 0x38, 0x50,
            0x6d, 0x51, 0x57, 0x5f, 0x62, 0x67, 0x68, 0x67, 0x3e, 0x4d, 0x71, 0x79, 0x70, 0x64,
            0x78, 0x5c, 0x65, 0x67, 0x63, 0xff, 0xdb, 0x00, 0x43, 0x01, 0x11, 0x12, 0x12, 0x18,
            0x15, 0x18, 0x2f, 0x1a, 0x1a, 0x2f, 0x63, 0x42, 0x38, 0x42, 0x63, 0x63, 0x63, 0x63,
            0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63,
            0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63,
            0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63, 0x63,
            0x63, 0x63, 0x63, 0x63, 0xff, 0xc0, 0x00, 0x11, 0x08, 0x00, 0x0b, 0x00, 0x0b, 0x03,
            0x01, 0x22, 0x00, 0x02, 0x11, 0x01, 0x03, 0x11, 0x01, 0xff, 0xc4, 0x00, 0x16, 0x00,
            0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x03, 0x01, 0x04, 0xff, 0xc4, 0x00, 0x21, 0x10, 0x00, 0x02, 0x01, 0x03,
            0x04, 0x03, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x02,
            0x03, 0x00, 0x11, 0x21, 0x04, 0x05, 0x12, 0x61, 0x14, 0x31, 0x41, 0x81, 0xff, 0xc4,
            0x00, 0x15, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0x05, 0xff, 0xc4, 0x00, 0x17, 0x11, 0x00, 0x03,
            0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x01, 0x11, 0x21, 0xff, 0xda, 0x00, 0x0c, 0x03, 0x01, 0x00, 0x02, 0x11, 0x03,
            0x11, 0x00, 0x3f, 0x00, 0x90, 0x6a, 0x23, 0x59, 0xa5, 0x59, 0x17, 0x92, 0xf1, 0x2e,
            0xcc, 0x4d, 0x88, 0xfb, 0x8e, 0xe9, 0x53, 0x69, 0xf2, 0x90, 0x4f, 0x2c, 0x73, 0x17,
            0x93, 0x26, 0xce, 0x6c, 0x3a, 0xfc, 0xf5, 0x59, 0x34, 0x0a, 0x1b, 0x77, 0x94, 0x30,
            0x04, 0x03, 0x6b, 0x1c, 0xe0, 0xda, 0xf4, 0x13, 0xea, 0x66, 0x59, 0x58, 0x2c, 0xac,
            0x00, 0xc0, 0x00, 0xd1, 0x4d, 0x88, 0xa4, 0xeb, 0xd3, 0xff, 0xd9,
        ];
        let dec = Decoder::new().unwrap();
        let image = dec.decode(LENA_JPG, PixelFormat::RgbInterleaved).unwrap();
        assert_eq!(image.dimensions(), (11, 11));
    }
}
